// csp.go (已修正)
package main

import (
	"math/big"
	"strconv"

	pbc "github.com/Nik-U/pbc"
)

// CloudServiceProvider 代表云服务商
type CloudServiceProvider struct {
	params     *SystemParams
	storedData map[string][]*pbc.Element // 存储文件块
	storedTags map[string][]*pbc.Element // 存储文件标签
}

// NewCSP 创建一个新的云服务商
func NewCSP(params *SystemParams) *CloudServiceProvider {
	return &CloudServiceProvider{
		params:     params,
		storedData: make(map[string][]*pbc.Element),
		storedTags: make(map[string][]*pbc.Element),
	}
}

// StoreFile 存储文件及其标签
func (csp *CloudServiceProvider) StoreFile(metadata *FileMetadata, data []*pbc.Element, tags []*pbc.Element) {
	csp.storedData[metadata.Name] = data
	csp.storedTags[metadata.Name] = tags
}

// ProofGen 生成数据持有证明
func (csp *CloudServiceProvider) ProofGen(metadata *FileMetadata, pk *PublicKey, tau string, c int) (*Proof, error) {
	blocks := csp.storedData[metadata.Name]
	tags := csp.storedTags[metadata.Name]

	// 1. 非交互式生成挑战
	challengeIndices := make(map[int]int) // map[index] -> count
	for i := 0; i < c; i++ {
		idx := csp.params.F(tau, i, metadata.NumBlocks)
		challengeIndices[idx]++
	}

	// 2. 聚合标签和数据
	sigmaAgg := csp.params.Pairing.NewG1().Set1()
	muPrime := csp.params.Pairing.NewZr().Set0()

	for j, count := range challengeIndices {
		vjData := []byte(tau + strconv.Itoa(j))
		vj := csp.params.Pairing.NewZr().SetFromHash(vjData)

		// 聚合标签: Π (σ_j)^vj
		sigma_j_vj := csp.params.Pairing.NewG1().PowZn(tags[j], vj)
		// 如果一个块被挑战多次, 乘以 count
		if count > 1 {
			countZr := csp.params.Pairing.NewZr().SetBig(big.NewInt(int64(count)))
			sigma_j_vj.PowZn(sigma_j_vj, countZr)
		}
		sigmaAgg.Mul(sigmaAgg, sigma_j_vj)

		// 聚合数据: Σ v_j * m_j
		countZr := csp.params.Pairing.NewZr().SetBig(big.NewInt(int64(count)))
		vjTotal := csp.params.Pairing.NewZr().Mul(vj, countZr)
		term := csp.params.Pairing.NewZr().Mul(vjTotal, blocks[j])
		muPrime.Add(muPrime, term)
	}

	// 3. 盲化处理
	s := csp.params.Pairing.NewZr().Rand()
	e_uv := csp.params.Pairing.NewGT().Pair(csp.params.U, pk.V)
	T := csp.params.Pairing.NewGT().PowZn(e_uv, s)
	gamma := csp.params.h(T)

	// μ = s + γ * μ'
	gamma_mu_prime := csp.params.Pairing.NewZr().Mul(gamma, muPrime)
	mu := csp.params.Pairing.NewZr().Add(s, gamma_mu_prime)

	return &Proof{
		Sigma: sigmaAgg,
		Mu:    mu,
		T:     T,
		Tau:   tau,
		C:     c,
	}, nil
}
